import numpy as np
import matplotlib.pyplot as plt
from init_control import UAVParameters
from uav_dynamics import UAVDynamics
from common_function import rotation_matrix
from PIDController import posion_pos_controller
from LQRController import LQRController
from square_signal import generate_ramp_signal

class UAVSimulation:
    def __init__(self, params, dynamics,position_controller,pose_controller):
        self.params = params
        self.dynamics = dynamics
        self.position_history = []
        self.setpoint_history = []
        self.position_controller = position_controller
        self.roll_controller = pose_controller[0]
        self.pitch_controller = pose_controller[1]
        self.yaw_controller = pose_controller[2]

    def motor_mixer(self, roll, pitch, yaw, thrust):
        idle_PWM = 1000
        scale = 1000
        M1 = (thrust - roll + pitch + yaw) * scale + idle_PWM
        M2 = (thrust + roll - pitch + yaw) * scale + idle_PWM
        M3 = (thrust + roll + pitch - yaw) * scale + idle_PWM
        M4 = (thrust - roll - pitch - yaw) * scale + idle_PWM
        return M1, M2, M3, M4
    def desired_euler_angles(self,eax, eay, psi):
        """
        Convert the desired horizontal acceleration to the desired Euler angle
        Input:
          eax, eay: desired horizontal acceleration
        Output:
          phi, theta: desired roll and pitch
        """
        g = 9.8  # 重力加速度 (m/s^2)

        phi = (-np.sin(psi) * eax + np.cos(psi) * eay) / g
        theta = (-np.cos(psi) * eax - np.sin(psi) * eay) / g

        return phi, theta
    def pwm_to_rpm(self, pwm):
        min_pwm = 1000
        max_pwm = 2000
        pwm = np.clip(pwm, min_pwm, max_pwm)
        rpm = np.clip((pwm - min_pwm) / (max_pwm - min_pwm), -1, 1)
        return rpm

    # def lqr_control(self, x, x_setpoint, A, B, Q, R):
    #     K, _, _ = control.lqr(A, B, Q, R)
    #     u = -K @ (x - x_setpoint)
    #     return u

    def run(self, square_corners, duration, dt):
        DEG2RAD = self.params.DEG2RAD
        posE = self.params.ModelInit_PosE.astype(np.float64)
        velB = self.params.ModelInit_VelB.astype(np.float64)
        velE = self.params.ModelInit_VelB.astype(np.float64)
        angEuler = self.params.ModelInit_AngEuler
        rate = self.params.ModelInit_RateB
        ModelParam_motorCr = self.params.ModelParam_motorCr
        ModelParam_motorWb = self.params.ModelParam_motorWb
        self.position_history.append(posE)

        state = [velB, angEuler, rate]
        time_steps = int(duration / dt)
        num_steps = time_steps // (len(square_corners) - 1)
        # next_corner_time = time_steps  // (len(square_corners)-1)
        input_signal = generate_ramp_signal(square_corners, num_steps)

        for step in range(time_steps):
            setpoint = input_signal[step]
            # 位置环LQR控制
            x_pos = np.hstack((posE, velE[:3]))
            x_setpoint_pos = np.hstack((setpoint, velE*dt))
            # u_pos = self.lqr_control(x_pos, x_setpoint_pos, A_pos, B_pos, Q_pos, R_pos)
            u_pos = self.position_controller.update(x_setpoint_pos,x_pos)
            # 姿态环LQR控制
            desired_phi, desired_theta = self.desired_euler_angles(u_pos[0], u_pos[1], angEuler[2])
            phi = np.clip(desired_phi, -35 * DEG2RAD, 35 * DEG2RAD)
            theta = np.clip(desired_theta, -35 * DEG2RAD, 35 * DEG2RAD)

            # x_att = np.hstack((angEuler, rate))
            # x_setpoint_att = np.array([desired_phi, desired_theta, 0, 0, 0, 0])
            # u_att = self.lqr_control(x_att, x_setpoint_att, A_att, B_att, Q_att, R_att)
            #
            # roll_output = u_att[0]
            # pitch_output = u_att[1]
            # yaw_output = u_att[2]
            roll_output = np.clip(self.roll_controller.update(outer_setpoint=phi, outer_measurement=angEuler[0],
                                                              inner_measurement=rate[0], dt=dt), -1, 1)
            pitch_output = np.clip(self.pitch_controller.update(outer_setpoint=theta, outer_measurement=angEuler[1],
                                                                inner_measurement=rate[1], dt=dt), -1, 1)
            yaw_output = np.clip(self.yaw_controller.update(outer_setpoint=0, outer_measurement=angEuler[2],
                                                            inner_measurement=rate[2], dt=dt), -1, 1)
            thrust = np.clip(-u_pos[2] + self.params.THR_HOVER, 0.05, 0.9)

            M1, M2, M3, M4 = self.motor_mixer(roll_output, pitch_output, yaw_output, thrust)

            w_target = (np.array(
                [self.pwm_to_rpm(M1), self.pwm_to_rpm(M2), self.pwm_to_rpm(M3), self.pwm_to_rpm(M4)]) * ModelParam_motorCr +
                        np.array([ModelParam_motorWb, ModelParam_motorWb, ModelParam_motorWb, ModelParam_motorWb]))

            state = self.dynamics.update_dynamics(state, w_target, dt)
            velB, angEuler, rate = state
            R = rotation_matrix(angEuler)
            velE = R @ velB
            posE = posE + velE * dt
            self.position_history.append(posE.copy())
            print(
                f"Time: {step * dt:.2f}s, Position: {posE}, Velocity_world: {velE}, Angle: {angEuler}, Rate: {rate}")
            self.setpoint_history.append(setpoint.copy())

    def plot_trajectory(self):
        position_history = np.array(self.position_history)
        setpoint_history = np.array(self.setpoint_history)
        fig = plt.figure()
        ax = fig.add_subplot(121, projection='3d')
        ax.plot(position_history[:, 0], position_history[:, 1], -position_history[:, 2], label='UAV Trajectory')
        ax.set_xlabel('X Position (m)')
        ax.set_ylabel('Y Position (m)')
        ax.set_zlabel('Z Position (m)')
        ax.set_title('Trajectory of UAV')
        ax.legend()
        ax = fig.add_subplot(122, projection='3d')
        ax.plot(setpoint_history[:, 0], setpoint_history[:, 1], -setpoint_history[:, 2], label='setpoint Trajectory')
        ax.set_xlabel('X Position (m)')
        ax.set_ylabel('Y Position (m)')
        ax.set_zlabel('Z Position (m)')
        ax.set_title('Desired Trajectory')
        ax.legend()

    def plot_state_error(self):
        self.position_controller.plot_error()
        self.roll_controller.outer_controller.plot_error('roll_angle')
        self.roll_controller.inner_controller.plot_error('roll_rate')
        self.pitch_controller.outer_controller.plot_error('pitch_angle')
        self.pitch_controller.inner_controller.plot_error('pitch_rate')
        self.yaw_controller.outer_controller.plot_error('yaw_angle')
        self.yaw_controller.inner_controller.plot_error('yaw_rate')
        # plt.show()

if __name__ == '__main__':
    params = UAVParameters()
    dynamics = UAVDynamics(params)
    # 定义位置环的状态空间模型
    A_pos = np.zeros((6, 6))
    A_pos[0, 3] = 1
    A_pos[1, 4] = 1
    A_pos[2, 5] = 1

    B_pos = np.zeros((6, 3))
    B_pos[3, 0] = 1
    B_pos[4, 1] = 1
    B_pos[5, 2] = 1

    Q_pos = np.eye(6) * 0.004
    R_pos = np.eye(3) * 0.004

    # 定义姿态环的状态空间模型
    A_att = np.zeros((6, 6))
    A_att[0, 3] = 1
    A_att[1, 4] = 1
    A_att[2, 5] = 1

    B_att = np.zeros((6, 3))
    B_att[3, 0] = 1
    B_att[4, 1] = 1
    B_att[5, 2] = 1

    Q_att = np.eye(6) * 0.05
    R_att = np.eye(3) * 2
    position_controller = LQRController(A_pos,B_pos,Q_pos,R_pos)
    pose_controller = posion_pos_controller(params).pos_controller()
    target_altitude = -5
    square_corners = [
        params.ModelInit_PosE.astype(np.float64),
        np.array([0, 0, target_altitude]),
        np.array([5, 0, target_altitude]),
        np.array([5, 5, target_altitude]),
        np.array([0, 5, target_altitude]),
        np.array([0, 0, target_altitude]),
        np.array([0, 0, 0])
    ]
    simulation = UAVSimulation(params, dynamics,position_controller,pose_controller)
    simulation.run(square_corners, 30.0, 0.01)
    simulation.plot_trajectory()
    simulation.plot_state_error()
    plt.show()